/*
 * Copyright (C) 2009, 2010, B3log Team
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.b3log.nide.util.parser.java.source;

import com.sun.source.tree.AnnotationTree;
import com.sun.source.tree.BinaryTree;
import java.util.logging.Level;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.type.DeclaredType;
import org.b3log.nide.util.parser.java.source.element.NideField;
import org.b3log.nide.util.parser.java.source.element.NideClass;
import com.sun.source.tree.ExpressionTree;
import java.lang.reflect.Field;
import javax.lang.model.element.Modifier;
import javax.lang.model.type.TypeKind;
import org.b3log.nide.util.parser.marker.IdentifierMarker;
import com.sun.source.tree.BlockTree;
import com.sun.source.tree.ClassTree;
import com.sun.source.tree.CompilationUnitTree;
import com.sun.source.tree.ErroneousTree;
import com.sun.source.tree.ExpressionStatementTree;
import com.sun.source.tree.IdentifierTree;
import com.sun.source.tree.IfTree;
import com.sun.source.tree.ImportTree;
import com.sun.source.tree.LineMap;
import com.sun.source.tree.LiteralTree;
import com.sun.source.tree.MemberSelectTree;
import com.sun.source.tree.MethodInvocationTree;
import com.sun.source.tree.MethodTree;
import com.sun.source.tree.ModifiersTree;
import com.sun.source.tree.NewArrayTree;
import com.sun.source.tree.NewClassTree;
import com.sun.source.tree.PrimitiveTypeTree;
import com.sun.source.tree.ReturnTree;
import com.sun.source.tree.Tree;
import com.sun.source.tree.VariableTree;
import com.sun.source.util.SourcePositions;
import com.sun.source.util.TreePath;
import com.sun.source.util.TreePathScanner;
import com.sun.source.util.Trees;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.logging.Logger;
import javax.lang.model.element.Element;
import javax.lang.model.util.Elements;
import org.b3log.nide.util.parser.Style;
import org.b3log.nide.util.parser.marker.HighlightingMarker;

/**
 * Class visitor.
 *
 * @author <a href="mailto:DL88250@gmail.com">Liang Ding</a>
 * @version 1.0.1.0, Nov 29, 2010
 */
public final class ClassVisitor extends TreePathScanner<Object, Void> {

    /**
     * Logger.
     */
    private static final Logger LOGGER =
            Logger.getLogger(ClassVisitor.class.getName());
    /**
     * Elements.
     */
    private Elements elements;
    /**
     * Imports.
     */
    private List<MemberSelectTree> imports = new ArrayList<MemberSelectTree>();
    /**
     * Source locations of identifiers.
     */
    private Set<IdentifierMarker> identifierMarkers =
            new HashSet<IdentifierMarker>();
    /**
     * Highlighting marker.
     */
    private Set<HighlightingMarker> highlightingMarkers =
            new HashSet<HighlightingMarker>();
    /**
     * Compilation unit tree.
     */
    private CompilationUnitTree compilationUnitTree;
    /**
     * Source positions.
     */
    private SourcePositions sourcePositions;
    /**
     * Line map.
     */
    private LineMap lineMap;
    /**
     * Java source code character sequence.
     */
    private CharSequence javaSourceCode;
    /**
     * Trees.
     */
    private Trees trees;
    /**
     * Variable trees.
     */
    private Map<Tree, Set<VariableTree>> variablesTrees =
            new HashMap<Tree, Set<VariableTree>>();
    /**
     * Keyword style.
     */
    private static final Style KEYWORD_STYLE = new Style();
    /**
     * Literal style.
     */
    private static final Style LITERAL_STYLE = new Style();
    /**
     * Class info.
     */
    private NideClass classInfo = new NideClass();

    static {
        KEYWORD_STYLE.setColor("#0000FF");
        LITERAL_STYLE.setColor("#0000FF");
    }

    /**
     * Constructs a class visitor for the specified elements and trees.
     *
     * @param elements the specified elements
     * @param trees trees
     */
    public ClassVisitor(final Elements elements, final Trees trees) {
        this.elements = elements;
        this.trees = trees;

        sourcePositions = trees.getSourcePositions();
    }

    @Override
    public Object visitImport(final ImportTree importTree, final Void param) {
        //LOGGER.finer("visitImport[" + importTree + "]");

        final Tree qualifiedIdentifierTree = importTree.getQualifiedIdentifier();
        final MemberSelectTree memberSelectTree =
                (MemberSelectTree) qualifiedIdentifierTree;

        imports.add(memberSelectTree);
        LOGGER.log(Level.FINER, "import[{0}.{1}]",
                   new Object[]{memberSelectTree.getExpression(),
                                memberSelectTree.getIdentifier()});

        final int startPosition = (int) sourcePositions.getStartPosition(
                compilationUnitTree, importTree);

        for (int i = 0; i < Keywords.IMPORT.length(); i++) {
            final HighlightingMarker highlightingMarker =
                    new HighlightingMarker(startPosition + i);
            highlightingMarker.setStyle(KEYWORD_STYLE);

            highlightingMarkers.add(highlightingMarker);
        }

        return super.visitImport(importTree, param);
    }

    @Override
    public Object visitCompilationUnit(
            final CompilationUnitTree compilationUnitTree, final Void param) {
        //LOGGER.finer("visitCompilationUnit[" + compilationUnitTree + "]");

        this.compilationUnitTree = compilationUnitTree;
        lineMap = compilationUnitTree.getLineMap();

        try {
            javaSourceCode = compilationUnitTree.getSourceFile().
                    getCharContent(false);

        } catch (final IOException e) {
            LOGGER.log(Level.SEVERE, e.getMessage(), e);
        }

        final ExpressionTree packageNameTree =
                compilationUnitTree.getPackageName();
        final long startPosition = sourcePositions.getStartPosition(
                compilationUnitTree, packageNameTree);
        final long lineNumber = lineMap.getLineNumber(startPosition);
        LOGGER.log(Level.FINER, "Package tree line number[{0}]", lineNumber);

        // FIXME: package org.b3log; package and expression tree MUST at the
        // same line and CAN NOT have comments between package and expression
        // tree.
        if (0 != lineNumber) {
            addHighlightingMarker(Keywords.PACKAGE.toString(),
                                  packageNameTree,
                                  -(Keywords.PACKAGE.length() + 1),
                                  KEYWORD_STYLE);
        }

        return super.visitCompilationUnit(compilationUnitTree, param);
    }

    @Override
    public Object visitClass(final ClassTree classTree, final Void param) {
        //LOGGER.finer("visitClass[" + classTree + "]");

        final ModifiersTree modifiersTree = classTree.getModifiers();

        final Set<Modifier> modifiers = modifiersTree.getFlags();
        int startPositionOffset = 0;
        for (final Modifier modifier : modifiers) {
            final String modifierString = modifier.toString();
            final Keywords keyword =
                    Keywords.valueOf(modifierString.toUpperCase());
            addHighlightingMarker(keyword.toString(), classTree,
                                  startPositionOffset,
                                  KEYWORD_STYLE);

            startPositionOffset += modifierString.length() + 1;
        }

        addHighlightingMarker(Keywords.CLASS.toString(), classTree,
                              startPositionOffset,
                              KEYWORD_STYLE); // TODO: interface
        startPositionOffset += Keywords.CLASS.length() + 1;

        final String classSimpleName = classTree.getSimpleName().toString();
        final Style classNameStyle = new Style();
        classNameStyle.setFontWeight("bold");
        addHighlightingMarker(classSimpleName, classTree,
                              startPositionOffset, classNameStyle);

        return super.visitClass(classTree, param);
    }

    @Override
    public Object visitBlock(final BlockTree blockTree, final Void param) {
        // LOGGER.finer("visitBlock[" + blockTree + "]");

        if (blockTree.isStatic()) {
            addHighlightingMarker(Keywords.STATIC.toString(),
                                  blockTree, 0, KEYWORD_STYLE);
        }

        variablesTrees.put(blockTree, new HashSet<VariableTree>());

        return super.visitBlock(blockTree, param);
    }

    @Override
    public Object visitExpressionStatement(
            final ExpressionStatementTree expressionStatementTree,
            final Void param) {
        //LOGGER.finer("visitExpressionStatement[" + expressionStatementTree + "]");

        return super.visitExpressionStatement(expressionStatementTree, param);
    }

    @Override
    public Object visitMemberSelect(final MemberSelectTree memberSelectTree,
                                    final Void param) {
//        LOGGER.finer("visitMemberSelect[" + memberSelectTree + "]");
//        final String identifier = memberSelectTree.getIdentifier().toString();
//        final int startPositionOffset =
//                memberSelectTree.getExpression().toString().length() + 1;
//
//        if (!identifier.endsWith(")")) {
//            final Style style = new Style();
//                style.setColor("green");
//
//
//            addHighlightingMarker(identifier, memberSelectTree,
//                                  startPositionOffset, style);
//        }

        return super.visitMemberSelect(memberSelectTree, param);
    }

    @Override
    public Object visitIdentifier(final IdentifierTree identifierTree,
                                  final Void param) {
        //LOGGER.finer("visitIdentifier[" + identifierTree + "]");
        final String identifierName = identifierTree.getName().toString();

        final Set<NideField> fieldInfos = classInfo.getFieldInfos();
        for (final NideField fieldInfo : fieldInfos) {
            if (fieldInfo.getName().equals(identifierName)) {
                final Style style = new Style();
                style.setColor("green");

                if (java.lang.reflect.Modifier.isStatic(
                        fieldInfo.getModifiers())) {
                    style.setFontStyle("italic");
                }

                addHighlightingMarker(identifierName,
                                      identifierTree, 0, style);
            }
        }

        return super.visitIdentifier(identifierTree, param);
    }

    @Override
    public Object visitNewArray(final NewArrayTree newArrayTree,
                                final Void param) {
        addHighlightingMarker(Keywords.NEW.toString(),
                              newArrayTree, 0, KEYWORD_STYLE);

        return super.visitNewArray(newArrayTree, param);
    }

    @Override
    public Object visitNewClass(final NewClassTree newClassTree,
                                final Void param) {
        addHighlightingMarker(Keywords.NEW.toString(),
                              newClassTree, 0, KEYWORD_STYLE);

        return super.visitNewClass(newClassTree, param);
    }

    @Override
    public Object visitLiteral(final LiteralTree literalTree, final Void param) {
        //LOGGER.finer("visitLiteral[" + literalTree + "]");

        try {
            final String literalString =
                    literalTree.toString().toUpperCase();
            final Literals literal = Literals.valueOf(literalString);

            addHighlightingMarker(literal.toString(),
                                  literalTree,
                                  0,
                                  LITERAL_STYLE);
        } catch (final IllegalArgumentException e) {
            LOGGER.fine(e.getMessage());
        }

        return super.visitLiteral(literalTree, param);
    }

    @Override
    public Object visitBinary(final BinaryTree binaryTree,
                              final Void param) {
        //LOGGER.finer("visitBinary[" + binaryTree + "]");

        return super.visitBinary(binaryTree, param);
    }

    @Override
    public Object visitReturn(final ReturnTree returnTree, final Void param) {
        addHighlightingMarker(Keywords.RETURN.toString(),
                              returnTree, 0, KEYWORD_STYLE);

        return super.visitReturn(returnTree, param);
    }

    @Override
    public Object visitIf(final IfTree ifTree, final Void param) {
        LOGGER.log(Level.FINER, "visitIf[{0}]", ifTree);

        return super.visitIf(ifTree, param);
    }

    @Override
    public Object visitMethodInvocation(
            final MethodInvocationTree methodInvocationTree, final Void param) {
        //LOGGER.finer("visitMethodInvocation[" + methodInvocationTree + "]");

        final TreePath currentPath = getCurrentPath();
        final Element element = trees.getElement(currentPath);
        //LOGGER.finer("Element[" + element + "]");


        return super.visitMethodInvocation(methodInvocationTree, param);
    }

    @Override
    public Object visitVariable(final VariableTree variableTree,
                                final Void param) {
        LOGGER.log(Level.FINER, "visitVariable[{0}]", variableTree);
        // TODO: nested block????

        final Tree parentTree = getCurrentPath().getParentPath().getLeaf();
        final Set<VariableTree> variableTrees = variablesTrees.get(parentTree);
        if (null != variableTrees) {
            variableTrees.add(variableTree);
        }

        // FIXME: All modifiers MUST at the same line and splitted by a whitespace,
        // may be we need a formatter to arrange the source code, such as annotation
        // @Inject format to @Inject()....
        int startPositionOffset = 0;
        final ModifiersTree modifiersTree = variableTree.getModifiers();
        final List<? extends AnnotationTree> annotations =
                modifiersTree.getAnnotations();
        for (final AnnotationTree annotationTree : annotations) {
            final String annotationString = annotationTree.toString();
            startPositionOffset += annotationString.length() + 1;
        }

        final Set<Modifier> modifiers = modifiersTree.getFlags();
        int mod = 0;

        for (final Modifier modifier : modifiers) {
            final String modifierString = modifier.toString();
            final Keywords keyword =
                    Keywords.valueOf(modifierString.toUpperCase());
            addHighlightingMarker(keyword.toString(), modifiersTree,
                                  startPositionOffset, KEYWORD_STYLE);

            startPositionOffset += modifierString.length() + 1;

            try {
                final Field field =
                        java.lang.reflect.Modifier.class.getField(modifierString.
                        toUpperCase());
                mod |= (Integer) field.get(java.lang.reflect.Modifier.class);
            } catch (final Exception e) {
                LOGGER.log(Level.SEVERE, e.getMessage(), e);
            }
        }

        final Tree tree = variableTree.getType();
        switch (tree.getKind()) {
            case PRIMITIVE_TYPE:
                final PrimitiveTypeTree primitiveTypeTree =
                        (PrimitiveTypeTree) tree;
                final TypeKind primitiveTypeKind =
                        primitiveTypeTree.getPrimitiveTypeKind();
                final Keywords keyword = Keywords.valueOf(
                        primitiveTypeKind.name());
                addHighlightingMarker(keyword.toString(), variableTree,
                                      startPositionOffset,
                                      KEYWORD_STYLE);
                startPositionOffset += keyword.length() + 1;
                break;
            case IDENTIFIER:
                final IdentifierTree identifierTree =
                        (IdentifierTree) tree;
                final String identifierName =
                        identifierTree.getName().toString();
                startPositionOffset += identifierName.length() + 1;
                break;
            default:
                LOGGER.log(Level.WARNING,
                           "Variable type tree[{0}] need to be handled", tree);
        }

        if (parentTree instanceof ClassTree) {
            final String variableName = variableTree.getName().toString();
            LOGGER.log(Level.FINER, "Found a field[name={0}]", variableName);
            final Style fieldNameStyle = new Style();
            fieldNameStyle.setColor("green");
            if (java.lang.reflect.Modifier.isStatic(mod)) {
                fieldNameStyle.setFontStyle("italic");
            }

            addHighlightingMarker(variableName, variableTree,
                                  startPositionOffset, fieldNameStyle);

            final NideField fieldInfo = new NideField(variableName,
                                                      mod);
            classInfo.addFieldInfo(fieldInfo);
            LOGGER.log(Level.FINE, "Added field info[{0}]", fieldInfo);
        }

        return super.visitVariable(variableTree, param);
    }

    @Override
    public Object visitMethod(final MethodTree methodTree, final Void param) {
        LOGGER.log(Level.FINER, "visitMethod[{0}]", methodTree);

        variablesTrees.put(methodTree, new HashSet<VariableTree>());
        final String methodName = methodTree.getName().toString();
        int startPositionOffset = 0;
        final Element element = trees.getElement(getCurrentPath());

        final List<? extends AnnotationMirror> annotationMirrors =
                element.getAnnotationMirrors();
        for (final AnnotationMirror annotationMirror : annotationMirrors) {
            final DeclaredType annotationType =
                    annotationMirror.getAnnotationType();
            // FIXME: annotation highlighting
            final String annotationSimpleName =
                    annotationType.asElement().getSimpleName().toString();
            LOGGER.log(Level.FINER,
                       "Method[name={0}] has annotation[simpleName={1}]",
                       new Object[]{methodName, annotationSimpleName});
            startPositionOffset +=
                    annotationSimpleName.length() + "@() ".length();
        }

        final ModifiersTree modifiersTree = methodTree.getModifiers();
        final Set<Modifier> modifiers = modifiersTree.getFlags();
        int modStartPosOffset = 0;
        for (final Modifier modifier : modifiers) {
            final String modifierString = modifier.toString();
            final Keywords keyword =
                    Keywords.valueOf(modifierString.toUpperCase());
            addHighlightingMarker(keyword.toString(), modifiersTree,
                                  modStartPosOffset, KEYWORD_STYLE);

            modStartPosOffset += modifierString.length() + 1;
        }

        startPositionOffset += modStartPosOffset;

        if ("<init>".equals(methodName)) {
            LOGGER.finer("<init>() method");

            // TODO: add constructor highlighting
            //addHighlightingMarkerForString(methodName, methodTree,
            //                               startPositionOffset);

            return super.visitMethod(methodTree, param);
        }

        final Tree returnTypeTree = methodTree.getReturnType();
        switch (returnTypeTree.getKind()) {
            case IDENTIFIER:
                final IdentifierTree returnTypeIdentifierTree =
                        (IdentifierTree) returnTypeTree;
                // TODO:
                LOGGER.log(Level.FINER, "Method return a type[{0}]",
                           returnTypeTree);
                startPositionOffset +=
                        returnTypeIdentifierTree.getName().length() + 1;
                break;
            case PRIMITIVE_TYPE:
                final PrimitiveTypeTree returnTypePrimitiveTypeTree =
                        (PrimitiveTypeTree) returnTypeTree;
                final String returnType =
                        returnTypePrimitiveTypeTree.getPrimitiveTypeKind().name();
                final Keywords keyword = Keywords.valueOf(returnType);
                addHighlightingMarker(keyword.toString(), methodTree,
                                      startPositionOffset,
                                      KEYWORD_STYLE);
                startPositionOffset += keyword.length() + 1;

                break;
            default:
                LOGGER.log(Level.WARNING,
                           "Method return type tree[{0}] need to be handled",
                           returnTypeTree);
        }

        final Style methodNameStyle = new Style();
        methodNameStyle.setColor("#000000");
        methodNameStyle.setFontWeight("bold");
        addHighlightingMarker(methodName, methodTree,
                              startPositionOffset,
                              methodNameStyle);

        return super.visitMethod(methodTree, param);
    }

    @Override
    public Object visitErroneous(final ErroneousTree erroneousTree,
                                 final Void param) {
        LOGGER.log(Level.FINER, "visitErroneous[{0}]", erroneousTree);
        final TreePath treePath = getCurrentPath();

        final List<? extends Tree> errorTrees = erroneousTree.getErrorTrees();
        for (final Tree errorTree : errorTrees) {
            if (errorTree instanceof IdentifierTree) {
                final IdentifierTree errorIdentifierTree =
                        (IdentifierTree) errorTree;
                final long startPosition =
                        sourcePositions.getStartPosition(
                        compilationUnitTree, erroneousTree);
                final long endPosition =
                        sourcePositions.getEndPosition(
                        compilationUnitTree, erroneousTree)
                        + errorIdentifierTree.getName().length();
                LOGGER.log(Level.FINER,
                           "Erroneous[identifierTree={0}, startPosition={1}, endPosition={2}]",
                           new Object[]{errorIdentifierTree,
                                        startPosition,
                                        endPosition});

                final CharSequence identifier = javaSourceCode.subSequence(
                        (int) startPosition, (int) endPosition);

                // TODO: nested block????
                final Tree parentTree = treePath.getParentPath().getParentPath().
                        getLeaf();
                final Set<VariableTree> variableTrees =
                        variablesTrees.get(parentTree);
                for (final VariableTree variableTree : variableTrees) {
                    if (variableTree.getName().equals(
                            errorIdentifierTree.getName())) {
                        final IdentifierMarker identifierMarker =
                                new IdentifierMarker((int) endPosition);
                        identifierMarker.setIdentifierName(identifier.toString());
                        identifierMarker.setIdentifierType(variableTree.getType().
                                toString());

                        identifierMarkers.add(identifierMarker);
                    }
                }
            }
        }

        return super.visitErroneous(erroneousTree, param);
    }

    /**
     * Gets imports.
     *
     * @return imports
     */
    public List<MemberSelectTree> getImports() {
        return Collections.unmodifiableList(imports);
    }

    /**
     * Gets the source markers of identifiers.
     *
     * @return source markers of identifiers
     */
    public Set<IdentifierMarker> getIdentifierMarkers() {
        return Collections.unmodifiableSet(identifierMarkers);
    }

    /**
     * Gets the source markers of syntax highlighting.
     *
     * @return source markers of syntax highlighting
     */
    public Set<HighlightingMarker> getHighlightingMarkers() {
        return Collections.unmodifiableSet(highlightingMarkers);
    }

    /**
     * Adds highlighting marker for the specified string, tree, the offset of
     * the string and style.
     *
     * @param string the specified string
     * @param tree the specified tree
     * @param startPositionOffset start position offset of the specified
     *         string
     * @param style the specified style
     */
    private void addHighlightingMarker(final String string,
                                       final Tree tree,
                                       final int startPositionOffset,
                                       final Style style) {
        final int startPosition = (int) sourcePositions.getStartPosition(
                compilationUnitTree, tree);

        for (int i = 0; i < string.length(); i++) {
            final HighlightingMarker stringHighlightingMarker =
                    new HighlightingMarker(startPosition + startPositionOffset
                                           + i);
            stringHighlightingMarker.setStyle(style);

            highlightingMarkers.add(stringHighlightingMarker);
        }
    }
}
